#!/usr/bin/python
# -*- coding:utf-8 -*-
# @FileName : DL5_test1_1.py
# Author    : myh

import torch
from torch import nn
from torch.nn import functional as F


class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = nn.Linear(20,256)
        self.lin2 = nn.Linear(256,256)
        self.out = nn.Linear(256,10)

    def forward(self,X):
        H1 = self.lin1(F.relu(X))
        H2 = self.lin2(F.relu(H1))
        out = self.out(F.relu(H2))
        return out

net = MLP()
print(net.weight[:])
X = torch.rand(size=(2, 4))
# print(net(X))
#
# print(net[2].state_dict())
#
# print(type(net[2].bias))
# print(net[2].bias)
# print(net[2].bias.data)
#
# print(*[(name, param.shape) for name, param in net[0].named_parameters()])
# print(*[(name, param.shape) for name, param in net.named_parameters()])


def block1():
    return nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                         nn.Linear(8, 4), nn.ReLU())

def block2():
    net = nn.Sequential()
    for i in range(4):
        # 在这里嵌套
        net.add_module(f'block {i}', block1())
    return net

# 初始化为正态分布
def init_normal(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.zeros_(m.bias)

# 初始化为1
def init_constant(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 1)
        nn.init.zeros_(m.bias)


def my_init(m):
    if type(m) == nn.Linear:
        print("Init", *[(name, param.shape)
                        for name, param in m.named_parameters()][0])
        nn.init.uniform_(m.weight, -10, 10)

        c,k = m.weight.data.shape
        for index1 in range(c):
            for index2 in range(k):
                val = m.weight.data[index1,index2]
                if -5<= val <= 5:
                    m.weight.data[index1,index2] = 0

        # m.weight.data *= m.weight.data.abs() >= 5


# print("before")
# print(net[0].weight[:])
# net.apply(my_init)
# print("after")
# print(net[0].weight[:])

# print(net[0].weight.data.shape)
# net.apply(init_constant)
# print(net[0].weight.data[0], net[0].bias.data[0])
#
#
# net.apply(init_normal)
# print(net[0].weight.data[0], net[0].bias.data[0])

# rgnet = nn.Sequential(block2(), nn.Linear(4, 1))
# print(rgnet(X))
# print(*[(name, param.shape) for name, param in rgnet.named_parameters()])

# 我们需要给共享层一个名称，以便可以引用它的参数
shared = nn.Linear(8, 8)
net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                    shared, nn.ReLU(),
                    shared, nn.ReLU(),
                    nn.Linear(8, 1))
net(X)
# 检查参数是否相同
print(net[2].weight.data[0] == net[4].weight.data[0])
net[2].weight.data[0, 0] = 100
# 确保它们实际上是同一个对象，而不只是有相同的值
print(net[2].weight.data[0] == net[4].weight.data[0])

